import unittest
import torch
import math
import matplotlib.pyplot as plt
from matplotlib.pyplot import plot, show

def as_tensor(arg):
    if isinstance(arg, torch.Tensor):
        return arg
    else:
        if isinstance(arg, complex):
            dtype = torch.complex
        elif hasattr(arg, 'dtype'):
            dtype = arg.dtype
        else:
            dtype = torch.double
        return torch.tensor(arg, dtype=dtype)

def wrap_tensor(func):
    return lambda *args: func(*map(as_tensor, args))

def scale(func, scale=1.0):
    return wrap_tensor(lambda t: func(t / as_tensor(scale)))

def translate(func, delay=0.0):
    return wrap_tensor(lambda t: func(t - as_tensor(delay)))

def amplify(func, amplitude=1.0):
    return wrap_tensor(lambda t: amplitude * func(t))

def line(time, center, bias, ramp_rate, width):
    bound = torch.logical_and(time < center + width/2, time > center - width/2) 
    value = bias + (time - center) * ramp_rate
    return bound * value

def rect(time, eps=1e-6):
    left = line(time, -0.5, 0.5, 1.0 / eps, eps)
    right = line(time, 0.5, 0.5, -1.0 / eps, eps)
    inner = line(time, 0.0, 1.0, 0.0, 1.0 - eps)
    return left + right + inner
    # left_bound = torch.logical_and(time < -0.5 + eps/2, time > -0.5 - eps/2) 
    # left_value = 0.5 + (time + 0.5) / eps
    # right_bound = torch.logical_and(time < 0.5 + eps/2, time > 0.5 - eps/2) 
    # right_value = 0.5 - (time - 0.5) / eps
    # inner_bound = torch.logical_and(time <= 0.5 - eps / 2, time >= -0.5 + eps/2)
    # inner_value = torch.ones_like(time)
    # return left_bound * left_value + inner_bound * inner_value + right_bound * right_value

def chirp(time, Tp, k, eps=1e-6):
    envelope = scale(rect, Tp)
    phase = k * torch.square(time) / 2.0
    return envelope(time) * torch.exp(2j * math.pi * phase)

class TestScale(unittest.TestCase):
    def testAsTensor(self):
        self.assertIsInstance(as_tensor(1), torch.Tensor, msg='as_tensor return non torch.Tensor')

    def testWrapTensor(self):
        f = wrap_tensor(torch.square)
        self.assertEqual(f(0.0), 0.0)
        self.assertIsInstance(f(1), torch.Tensor, msg='torch.square return non torch.Tensor')

    def testScale(self):
        g = scale(torch.square, 1.0)
        self.assertEqual(g(0), 0, msg='g(0) should be 0')

    def testTranslate(self):
        g = translate(torch.square, 1.0)
        self.assertEqual(g(0), 1.0)

    def testAmplify(self):
        g = amplify(torch.square, 2.0)
        self.assertEqual(g(1.0), 2.0)

    def testChirp(self):
        t = torch.linspace(-2, 2, 1000)
        data = chirp(t, 1.0, 100)
        plot(t, data.real, t, data.imag)
        show()


if __name__ == '__main__':
    plt.ion()
    unittest.main()
